package org.lzfto.flink.demo.business.flinkTest.source.mqtt;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

import org.apache.flink.api.connector.source.ReaderOutput;
import org.apache.flink.api.connector.source.SourceReader;
import org.apache.flink.core.io.InputStatus;
import org.eclipse.paho.client.mqttv3.IMqttDeliveryToken;
import org.eclipse.paho.client.mqttv3.MqttCallback;
import org.eclipse.paho.client.mqttv3.MqttClient;
import org.eclipse.paho.client.mqttv3.MqttConnectOptions;
import org.eclipse.paho.client.mqttv3.MqttException;
import org.eclipse.paho.client.mqttv3.MqttMessage;
import org.lzfto.flink.demo.config.mqtt.CustomMqttMessage;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.alibaba.fastjson2.JSON;
import com.alibaba.fastjson2.JSONObject;

public class MqttSourceReader implements SourceReader<String, MqttSplit> {

    private static final Logger LOGGER = LoggerFactory.getLogger(MqttSourceReader.class);

    private final MqttConnectionConfig config;
    private MqttClient mqttClient;
    private final BlockingQueue<String> messageQueue = new LinkedBlockingQueue<>(10000);
    private final List<MqttSplit> assignedSplits = Collections.synchronizedList(new ArrayList<>());
    private volatile boolean noMoreSplits = false;
    private CompletableFuture<Void> availableFuture = new CompletableFuture<>();
    private Integer taskQty;
    private Integer pollQty;
    

    public MqttSourceReader(MqttConnectionConfig config) {
        this.config = config;
        this.taskQty=0;
        this.pollQty=0;
        LOGGER.info("创建MqttSourceReader(clientId={},websocketClientName={},topic={})",config.getMqttClientId(),config.getWebsocketClientName(),
        		config.getTopic()
        		);
    }

    @Override
    public void start() {
        try {
        	//需要合理分配clientId，作为演示，仅仅设置一个固定的clientId
        	//其次同个主机上运行的时候，应该注意topic的不同，否则会异常
        	String clientId="MQTT-SOURCE-"+UUID.randomUUID().toString();
            mqttClient = new MqttClient(config.getArtemisHost(),clientId);
            MqttConnectOptions connOpts = new MqttConnectOptions();
            connOpts.setUserName(config.getArtemisUserName());
            connOpts.setPassword(config.getArtemisPassword().toCharArray());
            connOpts.setCleanSession(true);
            connOpts.setAutomaticReconnect(true);
            mqttClient.setCallback(new MqttCallback() {
                @Override
                public void connectionLost(Throwable cause) {
                    LOGGER.warn("MQTT 连接丢失，将自动重连", cause);
                }

                @Override
                public void messageArrived(String topic, MqttMessage message) {
                	JSONObject jo=JSONObject.parseObject(new String(message.getPayload()));
                	String pureDataMessage=null;
                	if (jo!=null) {
                		int flag=jo.getInteger("flag").intValue();
                		if (flag== CustomMqttMessage.MQTT_MESSAGE_TYPE_NORMAL) {
                			Object data= jo.get("data");
                			pureDataMessage=JSON.toJSONString(data);
                            taskQty++;
                		}
                		else {
                			if (flag==CustomMqttMessage.MQTT_MESSAGE_TYPE_START) {        				
                				taskQty=0;
                				pollQty=0;
                				LOGGER.info("开始处理订单数据");
                			}
                			else if (flag==CustomMqttMessage.MQTT_MESSAGE_TYPE_END) {
                				LOGGER.info("本次处理{}条订单数据",taskQty);
                			}
                		}
                	}
                	if (pureDataMessage!=null) {
                		boolean isOk=messageQueue.offer(pureDataMessage);
                    	if (isOk) {
                    		synchronized (this) {
                    	        availableFuture.complete(null);
                    	    }
                    	} else {
                    		LOGGER.warn("消息队列已满，丢弃消息");
                    	}
                	}
                }

                @Override
                public void deliveryComplete(IMqttDeliveryToken token) {
                    // 不需要处理
                }
            });

            mqttClient.connect(connOpts);
            mqttClient.subscribe(config.getTopic(), config.getQos());
        } catch (MqttException e) {
            throw new RuntimeException("初始化MQTT客户端失败", e);
        }
    }

    @Override
    public InputStatus pollNext(ReaderOutput<String> output) throws Exception {
    	String message = messageQueue.poll(100, TimeUnit.MILLISECONDS);
    	pollQty++;
        if (message != null) {
            output.collect(message);
            synchronized (this) {
            	availableFuture = new CompletableFuture<>();	
            }
            return InputStatus.MORE_AVAILABLE;
        }
        return noMoreSplits && assignedSplits.isEmpty()
                ? InputStatus.END_OF_INPUT
                : InputStatus.NOTHING_AVAILABLE;
    }

    @Override
    public CompletableFuture<Void> isAvailable() {
    	return availableFuture;
    }

    @Override
    public void addSplits(List<MqttSplit> splits) {
        if (splits == null || splits.isEmpty()) {
            LOGGER.warn("未分配新的 splits");
        } else {
            assignedSplits.addAll(splits);
        }
    }

    @Override
    public void notifyNoMoreSplits() {
        noMoreSplits = true;
        LOGGER.info("接收到 noMoreSplits 通知,不会有新的分片来了！");
    }

    @Override
    public List<MqttSplit> snapshotState(long checkpointId) {
        return new ArrayList<>(assignedSplits); // 可考虑记录位点等信息
    }

    @Override
    public void close() {
        if (mqttClient != null) {
            try {
                if (mqttClient.isConnected()) {
                    mqttClient.unsubscribe(config.getTopic());
                    mqttClient.disconnect();
                }
                mqttClient.close();
                LOGGER.info("在Mqtt Source中关闭 MQTT 客户端(clientId={},websocketClientName={})",config.getMqttClientId(),config.getWebsocketClientName());
            } catch (MqttException e) {
                LOGGER.warn("关闭 MQTT 客户端时发生异常", e);
            }
        }
    }
}
